import torch
import numpy as np


class LoadWeights(object):
    def __init__(self, path, device):
        super(LoadWeights, self).__init__()
        self.pretrain_weights_path = path
        self.device = device

    def load_weights(self):
        # 判断类型
        end_char = self.pretrain_weights_path.split(".")[-1]
        if end_char == "pth":
            pass
        if end_char == "npz":
            self.load_npz_weights()

    def load_torch_weights(self, model):
        model_dict = model.state_dict()
        pretrained_dict = torch.load(self.pretrain_weights_path, map_location=self.device)
        load_key, no_load_key, temp_dict = [], [], {}
        for k, v in pretrained_dict.items():
            if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
                temp_dict[k] = v
                load_key.append(k)
            else:
                no_load_key.append(k)
        model_dict.update(temp_dict)
        model.load_state_dict(model_dict)
        return model, no_load_key

    def load_npz_weights(self):
        from numpy import load
        file = load(self.pretrain_weights_path)
        # TODO 处理npz权重

